
import argparse
import os
import json
import cv2
from tqdm import tqdm
import pycocotools.mask as maskUtils
import numpy as np
import cv2
import pysrt
import re
import datetime
from resize import resize_ss
import shutil


def main():
    parser = argparse.ArgumentParser(description="Convert DJI Video&SRT to spire annotation")
    parser.add_argument(
        "--dji-dir",
        default="D:/dataset/240119-LH2300",
        help="Path to DJI Video&SRT dir",
    )
    parser.add_argument(
        "--output-dir",
        default="D:/dataset/BB-LH2300-v240216",
        help="Path to spire output dir",
    )
    parser.add_argument(
        "--skip",
        type=int,
        default=10,
        help="Number of frame to skip",
    )
    args = parser.parse_args()

    output_img_dir = os.path.join(args.output_dir, 'scaled_images')
    output_ann_dir = os.path.join(args.output_dir, 'annotations')
    if os.path.exists(output_img_dir):
        shutil.rmtree(output_img_dir)
    if os.path.exists(output_ann_dir):
        shutil.rmtree(output_ann_dir)

    if not os.path.exists(output_img_dir):
        os.makedirs(output_img_dir)
    if not os.path.exists(output_ann_dir):
        os.makedirs(output_ann_dir)

    n_skip = args.skip

    vid_names = os.listdir(args.dji_dir)
    for vid_name in vid_names:
        if vid_name.endswith('MP4'):
            name = os.path.splitext(vid_name)[0]
            vid_fn = os.path.join(args.dji_dir, vid_name)
            srt_fn = os.path.join(args.dji_dir, name + '.SRT')
            if os.path.exists(srt_fn):
                cap = cv2.VideoCapture(vid_fn)
                cap_len = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                ret, frame = cap.read()

                subs = pysrt.open(srt_fn)
                for id, sub in enumerate(subs):
                    start_time = sub.start.to_time()
                    end_time = sub.end.to_time()
                    ret, frame = cap.read()
                    if id % n_skip != 0:
                        continue
                    frame = resize_ss(frame, short_size=1080, max_size=1920, division=1)

                    current_date = datetime.date.today()
                    output_img_name = "{}_{:02}{:02}{:02}_{:06}.jpg".format(
                        name, current_date.year, current_date.month, current_date.day, sub.index)
                    spire_dict = {}
                    spire_dict['file_name'] = output_img_name
                    spire_dict['height'], spire_dict['width'] = frame.shape[0], frame.shape[1]
                    spire_dict['annos'] = []

                    text = sub.text
                    pattern = r"\[(.*?)\]"
                    results = re.findall(pattern, text)
                    for result in results:
                        pattern = r"(\w+) *: *([^ ,]+)"
                        results2 = re.findall(pattern, result)
                        for result2 in results2:
                            if result2[0] == 'Drone':
                                result3 = result2[1].split(':')
                                if result3[0] == 'Yaw':
                                    spire_dict['dji_yaw'] = float(result3[1])
                            elif result2[0] == 'iso':
                                spire_dict['dji_iso'] = float(result2[1])
                            elif result2[0] == 'shutter':
                                spire_dict['dji_shutter'] = result2[1]
                            elif result2[0] == 'fnum':
                                spire_dict['dji_fnum'] = float(result2[1])
                            elif result2[0] == 'ev':
                                spire_dict['dji_ev'] = float(result2[1])
                            elif result2[0] == 'ct':
                                spire_dict['dji_ct'] = float(result2[1])
                            elif result2[0] == 'color_md':
                                spire_dict['dji_color_md'] = result2[1]
                            elif result2[0] == 'focal_len':
                                spire_dict['dji_focal_len'] = float(result2[1])
                            elif result2[0] == 'latitude':
                                spire_dict['dji_latitude'] = float(result2[1])
                            elif result2[0] == 'longtitude':
                                spire_dict['dji_longtitude'] = float(result2[1])
                            elif result2[0] == 'rel_alt':
                                spire_dict['dji_rel_alt'] = float(result2[1])
                            elif result2[0] == 'abs_alt':
                                spire_dict['dji_abs_alt'] = float(result2[1])
                            elif result2[0] == 'Pitch':
                                spire_dict['dji_pitch'] = float(result2[1])
                            elif result2[0] == 'Roll':
                                spire_dict['dji_roll'] = float(result2[1])

                    print("{}, Frame: {}, Shape: {}".format(vid_name, sub.index, frame.shape))
                    cv2.imwrite(os.path.join(output_img_dir, output_img_name), frame)
                    output_fn = os.path.join(output_ann_dir, output_img_name + '.json')
                    with open(output_fn, "w") as f:
                        json.dump(spire_dict, f)
        print('pass')

    print('Done!')


if __name__ == '__main__':
    main()
